from torch.utils.data import DataLoader
import torch
import torchvision
import matplotlib.pyplot as plt
import numpy as np
from torchvision import transforms, datasets
import os

# 定义数据转换操作
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载MNIST测试数据集
testset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = DataLoader(testset, batch_size=32, shuffle=False)


# 创建图像保存目录
save_dir = 'img'
os.makedirs(save_dir, exist_ok=True)

# 打印图像函数
def imshow(img, filepath=None):
    img = img / 2 + 0.5  # 去归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)), cmap='gray')
    if filepath:
        plt.savefig(filepath)  # 保存图像
    plt.show()

# 保存前20个图像
dataiter = iter(testloader)
for i in range(20):
    images, labels = next(dataiter)
    # 获取第i张图片
    img = images[i % len(images)].unsqueeze(0) # 保证不会越界
    # 保存图片
    imshow(torchvision.utils.make_grid(img), filepath=os.path.join(save_dir, f"{i+1}.jpg"))